import java.io.*
import java.lang.IllegalArgumentException
import java.lang.StringBuilder
import java.security.MessageDigest
import java.util.*
import java.util.zip.ZipFile
import kotlin.system.exitProcess

fun Run(inCmd: List<String>, inWorkdir: String? = null) {
    println("Run: $inCmd")
    val pb: ProcessBuilder = ProcessBuilder(inCmd)
            .directory(File(inWorkdir ?: "."))
            .redirectErrorStream(true)
    val p = pb.start()
    val reader = BufferedReader(InputStreamReader(p.inputStream))
    while (true) {
        val line = reader.readLine() ?: break
        println(line)
    }
    p.waitFor()
    if (0 != p.exitValue()) {
        throw RuntimeException()
    }
}

fun CheckOutput(inCmd: List<String>, inWorkdir: String? = null): List<String> {
    println("Run: $inCmd")
    val ret = mutableListOf<String>()
    val pb: ProcessBuilder = ProcessBuilder(inCmd)
            .directory(File(inWorkdir ?: "."))
            .redirectErrorStream(true)
    val p = pb.start()
    val reader = BufferedReader(InputStreamReader(p.inputStream))
    while (true) {
        val line = reader.readLine() ?: break
        ret.add(line)
    }
    p.waitFor()
    if (0 != p.exitValue()) {
        throw RuntimeException()
    }
    return ret
}

fun getImgInfo(zipFIle: File, name: String): List<Any> {
    println("get image [$name] from zip file")
    val ret = mutableListOf<Any>()
    val messageDigest = MessageDigest.getInstance("SHA1")
    val zf = ZipFile(zipFIle)
    val entry = zf.getEntry(name) ?: throw IllegalArgumentException("$name not in ${zipFIle.path}")
    val iS = zf.getInputStream(entry)
    val buffer = ByteArray(1024)
    while (true) {
        val read = iS.read(buffer)
        if (read == -1) {
            break
        }
        messageDigest.update(buffer, 0, read)
    }
    val formatter = Formatter()
    messageDigest.digest().forEach {
        formatter.format("%02x", it)
    }
    ret.add(name.replace(".", "_"))
    ret.add(entry.size)
    ret.add(formatter.toString())
    return ret
}

fun probeParts() {
    val parts = CheckOutput("adb shell ls /dev/block/by-name".split(" "))
    println("Found Partitions: $parts")
    val abParts = mutableListOf<String>()
    val monoParts = mutableListOf<String>()
    val tmpParts = mutableListOf<String>()
    for (part in parts) {
        if (part.endsWith("_a") || part.endsWith("_b")) {
            val stem = part.substring(0, part.length - 2)
            if (tmpParts.contains(stem)) {
                tmpParts.remove(stem)
                abParts.add(stem)
            } else {
                tmpParts.add(stem)
            }
        } else {
            monoParts.add(part)
        }
    }
    if (tmpParts.size > 0) {
        println("ERROR: can not parse partitions : $tmpParts")
    }
    println("AB partitions  : $abParts")
    println("mono partitions: $monoParts")

    FileOutputStream("parts.txt").use { fos ->
        val props = Properties()
        abParts.forEach {
            props.setProperty("ab.$it.img", "/dev/block/by-name/$it")
        }
        monoParts.forEach {
            props.setProperty("$it.img", "/dev/block/by-name/$it")
        }
        props.store(fos, "probed partitions")
    }
}

fun verifyParts() {
    val scriptName = "verify.sh"
    val signedZip = "signed.zip"
    val props = Properties()
    props.load(FileInputStream("parts.txt"))
    println(props)

    val dollar = "$"
    val scriptTemplate = """if ! applypatch --check EMMC:%s:%d:%s >/dev/null ; then
    %s=" X "
    result="$dollar{result}1"
else
    %s=" - "
    result="$dollar{result}0"
fi"""
    val pw = PrintWriter(File(scriptName))
    val abSB = StringBuilder()
    val monoSB = StringBuilder()

    val abTemp = """print "%-20s: $dollar%s     $dollar%s""""
    val monoTemp = """print "%-20s:    $dollar%s""""
    props.forEach { t, u ->
        if (t.toString().startsWith("ab.")) {
            val info = getImgInfo(File(signedZip), t.toString().substring(3))
            pw.println(String.format(scriptTemplate, u.toString() + "_a", info[1], info[2], info[0].toString() + "_a", info[0].toString() + "_a"))
            pw.println(String.format(scriptTemplate, u.toString() + "_b", info[1], info[2], info[0].toString() + "_b", info[0].toString() + "_b"))
            abSB.appendln(String.format(abTemp, t.toString().substring(3), info[0].toString() + "_a", info[0].toString() + "_b"))
        } else {
            val info = getImgInfo(File(signedZip), t.toString())
            pw.println(String.format(scriptTemplate, u, info[1], info[2], info[0], info[0]));
            monoSB.appendln(String.format(monoTemp, t.toString(), info[0]))
        }
    }
    pw.println(String.format(monoTemp, "    mono", "x result"))
    pw.println("echo '--------------------------------'")
    pw.println(monoSB.toString())
    pw.println("echo")
    pw.println(String.format(abTemp, "    A/B", "x A", "x  B"))
    pw.println("echo '--------------------------------'")
    pw.println(abSB.toString())
    pw.close()
}

fun usage() {
    println("Usage: <prog> probe")
    println("Usage: <prog> verify")
}

if (args.size != 1) {
    usage()
    exitProcess(1)
}

when (args[0]) {
    "probe" -> {
        probeParts()
    }
    "verify" -> {
        verifyParts()
    }
    else -> {
        usage()
        exitProcess(1)
    }
}
